# Social Mobilisation in Partisan Spaces
# Replication Materials

# Functions to create custom plots from interflex::inter.binning() output
# last change: 2020-04-09

library("ggplot2")
loadNamespace("dplyr")

# This code is heavily lifted from interflex::inter.binning() and
# interflex::inter.plot() (see below for details).
# Everything is based on the code in version 1.0.8 of interflex.


# plot_binning() -----------------------------------------------------

# This is a slightly adapted version of interflex::inter.plot() which
# gives us more flexibility in manually adjusting the labels and other
# attributes of the graph.

plot_binning <- function(mod, include_linear = FALSE, adjusted_density = NULL) {

  # If adjusted density estimates are available, use these instead of the
  # ones in "mod"

  if (!is.null(adjusted_density)) {
    if (!is.list(adjusted_density)) {
      stop('"adjusted_density" needs to be a list')
    }
    # Check adjusted_density names/classes
    if (!identical(
      x = sort(c("de.tr", "de.co")),
      y = sort(names(adjusted_density))
    )) {
      stop('"adjusted_density" needs to contain two elements called "de.co" and "de.tr"')
    }
    if (!(class(adjusted_density$de.co) == "density" & class(adjusted_density$de.tr) == "density")) {
      stop('"de.co" and "de.tr" in "adjusted_density" need to be of class "density"')
    }

    # Replace density estimates
    mod$de.co <- adjusted_density$de.co
    mod$de.tr <- adjusted_density$de.tr
  }

  yrange <- c(mod$est.bin$CI_lower, mod$est.bin$CI_upper)
  maxdiff <- (max(yrange) - min(yrange))

  deX.ymin <- min(yrange) - maxdiff/5
  deX <- data.frame(
    x = mod$de.co$x,
    y = mod$de.co$y / max(mod$de.co$y) *  maxdiff/5 + min(yrange) - maxdiff/5
  )
  deX.tr <- data.frame(
    x =  mod$de.tr$x,
    y =  mod$de.tr$y / max( mod$de.tr$y) * maxdiff/5 + min(yrange) - maxdiff/5
  )

  # Set up plot
  plot_out <- ggplot(mod$est.bin, aes(x = x0)) +
    geom_hline(yintercept = 0, colour = "white", size = 2)

  # Include linear estimate if requested
  if (include_linear) {
    plot_out <- plot_out +
      geom_line(data = mod$est.lin, aes(X.lvls, marg)) +
      geom_ribbon(data = mod$est.lin, aes(x = X.lvls, ymin = lb,
                                      ymax = ub), alpha = 0.2)
  }

  # Add bars
  plot_out <- plot_out +
    geom_errorbar(aes(ymin = CI_lower, ymax = CI_upper), width = 0.1, colour = "red") +
    geom_point(aes(y = coef), size = 4, fill = "white", stroke = 1, shape = 21, colour = "red")

  # Add density
  plot_out <- plot_out +
    geom_ribbon(
      data = deX,
      aes(x = x, ymax = y, ymin = deX.ymin), alpha = 0.2
    ) +
    geom_ribbon(
      data = deX.tr,
      aes(x = x, ymax = y, ymin = deX.ymin),
      alpha = 0.2,
      fill = "red"
    )

  # Add text labels
  # plot_out <- plot_out +
  #   annotate("text",
  #            x = mod$est.bin$x0,
  #            y = max(yrange) + maxdiff/5,
  #            label = c("All Labour", "Mixed Partisan", "All Rival"),
  #            colour = "gray50",
  #            size = 7)

  # Fix scales and add axis labels
  plot_out <- plot_out +
    scale_x_continuous(breaks = c(0, 0.5, 1), labels = c("All Labour", "Mixed Partisan", "All Rival")) +
    labs(
      x = "Moderator: Partisan Composition of Closest Household",
      y = "Marginal Effect of Letter on Turnout"
    )

  return(plot_out)
}


# inter_beginning_density() ------------------------------------------

# The interflex::inter.binning() output includes kernel density estimates
# for the distribution the moderator for treatment and control groups.
# This density estimation uses the continuous version of the moderator even
# when supplying cutoffs to the function.

# In our case, it makes more sense to show kernel density estimates on the
# binned variable.

# This requires re-calculating the density separately as the raw data is not
# provided as output by interflex::inter.binning().

# This function requires "data", "cutoffs", "X", and "D" specified
# in the same way as in the original call to inter.binning()

inter_binning_density <- function(D, X, cutoffs, data, weights) {
  # Re-create cut-offs (copied from lines 298 to 301)
  cutoffs <- cutoffs[which(cutoffs > min(data[, X]) & cutoffs < max(data[, X]))]
  cuts.X <- sort(unique(c(min(data[, X]), cutoffs, max(data[, X]))))
  groupX <- cut(data[, X], breaks = cuts.X, labels = FALSE)
  groupX[which(data[, X] == min(data[, X]))] <- 1

  # "groupX" now contains the groups (and is ordered in the same way as "data").
  # Recode the this variable such that:
  # 1 -> 0
  # 2 -> 0.5
  # 3 -> 1
  X_groups  <- dplyr::recode(groupX, `1` = 0, `2` = 0.5, `3` = 1)

  # Add this variable to "data"
  data$X_groups <- X_groups

  # Re-calculate density with "X_grouped" instead of "X"
  X <- "X_groups"
  suppressWarnings(
    de.co <- density(data[data[, D] == 0, X], weights = data[data[, D] == 0, weights])
  )
  suppressWarnings(
    de.tr <- density(data[data[, D] == 1, X], weights = data[data[, D] == 1, weights])
  )

  # Return a list with the elements named as in the output from
  # inter.binning()
  list(
    de.co = de.co,
    de.tr = de.tr
  )
}
